// network/fc_layer.go 修复版本
package network

import (
	"gonum.org/v1/gonum/mat"
	"math"
	"math/rand"
)

// FCLayer 全连接层
type FCLayer struct {
	InputSize            int
	OutputSize           int
	FCWeights            *mat.Dense
	FCBiases             *mat.VecDense
	Activation           func(*mat.VecDense) *mat.VecDense
	ActivationDerivative func(*mat.VecDense) *mat.VecDense
	Inputs               *mat.VecDense // 保存前向传播的输入，用于反向传播
	PreActivation        *mat.VecDense // 保存激活前的值，用于反向传播
}

// NewFCLayer 创建一个新的全连接层
func NewFCLayer(inputSize, outputSize int, activation func(*mat.VecDense) *mat.VecDense, activationDerivative func(*mat.VecDense) *mat.VecDense) *FCLayer {
	weights := mat.NewDense(outputSize, inputSize, nil)
	scale := math.Sqrt(2.0 / float64(inputSize)) // He初始化
	for i := 0; i < outputSize; i++ {
		for j := 0; j < inputSize; j++ {
			weights.Set(i, j, rand.NormFloat64()*scale)
		}
	}
	biases := mat.NewVecDense(outputSize, nil)

	return &FCLayer{
		InputSize:            inputSize,
		OutputSize:           outputSize,
		FCWeights:            weights,
		FCBiases:             biases,
		Activation:           activation,
		ActivationDerivative: activationDerivative,
	}
}

// Forward 实现全连接层的前向传播
func (f *FCLayer) Forward(input interface{}) interface{} {
	inputVec, ok := input.(*mat.VecDense)
	if !ok {
		panic("FCLayer.Forward: input is not a *mat.VecDense")
	}

	// 保存输入用于反向传播
	f.Inputs = mat.VecDenseCopyOf(inputVec)

	// 计算前激活值 z = Wx + b
	z := mat.NewVecDense(f.OutputSize, nil)
	z.MulVec(f.FCWeights, inputVec)
	z.AddVec(z, f.FCBiases)

	// 保存前激活值用于反向传播
	f.PreActivation = mat.VecDenseCopyOf(z)

	// 应用激活函数
	return f.Activation(z)
}

// Backward 实现全连接层的反向传播
func (f *FCLayer) Backward(outputGradient interface{}, learningRate float64) interface{} {
	gradVec, ok := outputGradient.(*mat.VecDense)
	if !ok {
		panic("FCLayer.Backward: outputGradient is not a *mat.VecDense")
	}

	// 检查是否有前向传播的输入
	if f.Inputs == nil || f.PreActivation == nil {
		panic("FCLayer.Backward: Forward must be called before Backward")
	}

	var deltaVec *mat.VecDense

	// 检查是否有激活函数导数
	if f.ActivationDerivative != nil {
		// 计算对前激活值的梯度
		activationGrad := f.ActivationDerivative(f.PreActivation)
		deltaVec = mat.NewVecDense(f.OutputSize, nil)
		for i := 0; i < f.OutputSize; i++ {
			deltaVec.SetVec(i, gradVec.AtVec(i)*activationGrad.AtVec(i))
		}
	} else {
		// 如果没有激活函数导数（例如Softmax输出层），直接使用输出梯度
		deltaVec = gradVec
	}

	// 计算权重梯度
	weightGrad := mat.NewDense(f.OutputSize, f.InputSize, nil)
	for i := 0; i < f.OutputSize; i++ {
		for j := 0; j < f.InputSize; j++ {
			weightGrad.Set(i, j, deltaVec.AtVec(i)*f.Inputs.AtVec(j))
		}
	}

	// 计算输入梯度（用于传递给前一层）
	inputGrad := mat.NewVecDense(f.InputSize, nil)
	for i := 0; i < f.InputSize; i++ {
		sum := 0.0
		for j := 0; j < f.OutputSize; j++ {
			sum += f.FCWeights.At(j, i) * deltaVec.AtVec(j)
		}
		inputGrad.SetVec(i, sum)
	}

	// 更新参数
	for i := 0; i < f.OutputSize; i++ {
		// 更新偏置
		newBias := f.FCBiases.AtVec(i) - learningRate*deltaVec.AtVec(i)
		f.FCBiases.SetVec(i, newBias)

		// 更新权重
		for j := 0; j < f.InputSize; j++ {
			newWeight := f.FCWeights.At(i, j) - learningRate*weightGrad.At(i, j)
			f.FCWeights.Set(i, j, newWeight)
		}
	}

	return inputGrad
}

// GetType 返回层的类型
func (f *FCLayer) GetType() string {
	return "FullyConnected"
}

// GetOutputShape 返回输出形状
func (f *FCLayer) GetOutputShape() []int {
	return []int{f.OutputSize}
}
